#!/usr/bin/python

from opencv.cv import *
from opencv.highgui import *
from min_quad import min_quad

from IPython.Shell import IPythonShellEmbed
ipshell = IPythonShellEmbed([])

storage = None

def get_storage():
    global storage
    if storage is None:
        storage = cvCreateMemStorage(0)
    return storage

class ImageOp(object):
    def __init__(self):
        pass
    def __rshift__(self, next):
        "self >> next"
        self.run()
        next.input = self.output
        return next
    def run(self):
        raise NotImplementedError
    def set_input(self, input):
        raise NotImplementedError
    def get_output(self):
        raise NotImplementedError

class Identity(ImageOp):
    def __init__(self):
        pass
    def run(self):
        self.output = self.input

class End(ImageOp):
    def __init__(self):
        pass
    def run(self):
        self.output = None

class LoadImage(ImageOp):
    GRAYSCALE = 0
    COLOR = 1
    ANYDEPTH = 2
    ANYCOLOR = 4
    def __init__(self, fname, flags=COLOR):
        self.fname = fname
        self.flags = flags
    def run(self):
        self.output = cvLoadImage(self.fname, self.flags)

class FromImage(ImageOp):
    def __init__(self, image):
        self.image = image
    def run(self):
        self.output = self.image

class SaveImage(ImageOp):
    def __init__(self, fname):
        self.fname = fname
    def run(self):
        cvSaveImage(self.fname, self.input)
        self.output = self.input

class ShowImage(ImageOp):
    def __init__(self, window=None):
        self.window = window or ''
    def run(self):
        cvNamedWindow(self.window)
        cvShowImage(self.window, self.input)
        cvWaitKey(0)
        self.output = self.input

class ConvertBGR2GRAY(ImageOp):
    def __init__(self):
        pass
    def run(self):
        self.output = cvCreateImage(cvGetSize(self.input), 8, 1)
        cvCvtColor(self.input, self.output, CV_BGR2GRAY)

class ConvertGRAY2BGR(ImageOp):
    def __init__(self):
        pass
    def run(self):
        self.output = cvCreateImage(cvGetSize(self.input), 8, 3)
        cvCvtColor(self.input, self.output, CV_GRAY2BGR)

class Canny(ImageOp):
    def __init__(self, threshold1, threshold2, aperture_size=3):
        self.threshold1 = threshold1
        self.threshold2 = threshold2
        self.aperture_size = aperture_size
    def run(self):
        self.output = cvCreateImage(cvGetSize(self.input), 8, 1)
        cvCanny(self.input, self.output, self.threshold1, self.threshold2,
                self.aperture_size)

class Segmentation(ImageOp):
    def __init__(self, threshold1, threshold2, level=4):
        self.threshold1 = threshold1
        self.threshold2 = threshold2
        self.level = level
    def run(self):
        self.input.width &= -(1<<self.level)  # XXX: does not currently work
        self.input.height &= -(1<<self.level)  # XXX: does not currently work
        self.output = cvCloneImage(self.input)
        comp = cvPyrSegmentation(self.input, self.output, get_storage(),
                                 self.level, self.threshold1, self.threshold2)

class FindContours(ImageOp):
    def __init__(self):
        pass
    def run(self):
        self.output = cvCreateImage(cvGetSize(self.input), 8, 3)
        cvCvtColor(self.input, self.output, CV_GRAY2BGR)

        count, contours = cvFindContours(self.input, get_storage(),
                                         sizeof_CvContour,
                                         CV_RETR_EXTERNAL,
                                         CV_CHAIN_APPROX_SIMPLE,
                                         cvPoint(0,0))
        #print count, contours, cvContourPerimeter(contours)
        cvDrawContours(self.output, contours,
                       CV_RGB(255, 0, 0), CV_RGB(255, 0, 0), 1,
                       3, CV_AA, cvPoint(0,0))
        
        contours = cvApproxPoly(contours, sizeof_CvContour, get_storage(),
                                CV_POLY_APPROX_DP, 7, 1)
        cvDrawContours(self.output, contours,
                       CV_RGB(0, 255, 0), CV_RGB(255, 0, 0), 1,
                       2, CV_AA, cvPoint(0,0))

        # vertices = CvSeq_CvPoint.cast(
        #     cvCreateSeq(0, sizeof_CvSeq, sizeof_CvPoint, get_storage()))
        vertices = []

        for contour in contours.hrange():
            # contour = cvApproxPoly(contour, sizeof_CvContour, get_storage(),
            #                        CV_POLY_APPROX_DP, 7, 1)
            perimeter = cvArcLength(contour, cvSlice(0, 0x3fffffff), 0)
            area = abs(cvContourArea(contour ))
            # print perimeter, cvContourPerimeter(contour), area
            # TODO(nilton): fix hardcoded values below
            # change 200 to 0.33 * min(w, h)
            if (perimeter > 1 and area > 1) or perimeter > 200:
                cvDrawContours(self.output, contour,
                               CV_RGB(0, 0, 255), CV_RGB(0, 255, 0), 1,
                               1, CV_AA, cvPoint(0,0))
                for p in contour:
                    cvCircle(self.output, p, 2, CV_RGB(255, 255, 255),
                             CV_FILLED, CV_AA, 0)
                    vertices.append(p)

        hull = cvConvexHull2(vertices, CV_COUNTER_CLOCKWISE)
        #hpoly = cvCreateMat(1, len(hull), ?)
        hullpoly = cvCreateSeq(CV_SEQ_ELTYPE_POINT | CV_SEQ_CONTOUR,
                               sizeof_CvContour, sizeof_CvPoint, get_storage())
        hullpoly = CvSeq_CvPoint.cast(hullpoly)
        p0 = vertices[hull[-1]]
        hullpoly.append(p0)
        for i in hull:
            p1 = vertices[i]
            hullpoly.append(p1)
            cvLine(self.output, p0, p1, CV_RGB(255, 255, 0), 2, CV_AA, 0)
            p0 = p1

        poly = cvApproxPoly(hullpoly, hullpoly.header_size,
                            get_storage(), CV_POLY_APPROX_DP, 5, 0)
        # print hullpoly.total, poly.total
        cvDrawContours(self.output, poly, CV_RGB(0, 255, 255), CV_RGB(0, 0, 0),
                       1, 3, CV_AA, cvPoint(0,0))

        quad = min_quad([ (p.x, p.y) for p in hullpoly ])
        quad = [ cvPoint(x, y) for (x, y) in quad ]
        p0 = quad[-1]
        for p1 in quad:
            cvLine(self.output, p0, p1, CV_RGB(127,127,127), 2, CV_AA)
            p0 = p1

        quad.sort(key=lambda p: p.x)
        ul, ll, ur, lr = quad  # first guess
        if ul.y > ll.y:
            ul, ll = ll, ul
        if ur.y > lr.y:
            ur, lr = lr, ur
        self.points = tuple( (p.x, p.y) for p in (ul, ll, lr, ur) )




if __name__ == '__main__':
    import sys
    fname = sys.argv[1]
    #outname = sys.argv[2]

    # old parameters
    # Segmentation(200, 60)
    # Segmentation(75, 60)
    # Canny(50, 200)


    find_contours = FindContours()
    (
        LoadImage(fname) >>
        ShowImage('img') >>
        Segmentation(75, 65) >>
        ShowImage('img') >>
        ConvertBGR2GRAY() >>
        Canny(50, 180) >>
        ShowImage('img') >>
        find_contours >>
        ShowImage('img') >>
        #SaveImage('output.jpg') >> 
        End()
    )

    import Image
    from perspective_correction import perspective_correction
    im = perspective_correction(Image.open(fname), *find_contours.points)
    im.show()
